"""
etcd状态存储后端 - 强一致性状态存储

这个模块实现了基于etcd的状态存储：
- 强一致性保证
- 分布式锁支持
- 配置管理
- 事务支持
- 租约和TTL管理
"""

import asyncio
import json
import logging
import time
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, List, Optional, Set, Union
import etcd3

from .result_aggregator import StateStorage, TaskRecord, TaskStatus

logger = logging.getLogger(__name__)


class EtcdStateStorage(StateStorage):
    """etcd状态存储实现"""
    
    def __init__(self, 
                 host: str = "localhost",
                 port: int = 2379,
                 key_prefix: str = "/gpu_auto_scaling",
                 default_ttl: int = 86400 * 7,  # 7天
                 timeout: int = 5):
        self.host = host
        self.port = port
        self.key_prefix = key_prefix
        self.default_ttl = default_ttl
        self.timeout = timeout
        
        self.etcd_client: Optional[etcd3.Etcd3Client] = None
        
        # 键名模板
        self.task_key_template = f"{key_prefix}/task/{{task_id}}"
        self.task_status_key_template = f"{key_prefix}/status/{{status}}"
        self.server_tasks_key_template = f"{key_prefix}/server/{{server_id}}/tasks"
        self.task_index_key = f"{key_prefix}/task_index"
        self.lock_key_template = f"{key_prefix}/lock/{{resource}}"
        self.lease_key_template = f"{key_prefix}/lease/{{resource}}"
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("etcd状态存储已初始化")
    
    async def connect(self) -> None:
        """连接到etcd"""
        try:
            # 创建etcd客户端
            self.etcd_client = etcd3.client(
                host=self.host,
                port=self.port,
                timeout=self.timeout
            )
            
            # 测试连接
            self.etcd_client.status()
            
            self.logger.info("etcd连接成功")
            
        except Exception as e:
            self.logger.error(f"etcd连接失败: {e}")
            raise
    
    async def disconnect(self) -> None:
        """断开etcd连接"""
        if self.etcd_client:
            self.etcd_client.close()
            self.etcd_client = None
        
        self.logger.info("etcd连接已断开")
    
    async def get_task(self, task_id: str) -> Optional[TaskRecord]:
        """获取任务记录"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            task_key = self.task_key_template.format(task_id=task_id)
            value, metadata = self.etcd_client.get(task_key)
            
            if not value:
                return None
            
            # 解析任务数据
            task_dict = json.loads(value.decode())
            return self._dict_to_task_record(task_dict)
            
        except Exception as e:
            self.logger.error(f"获取任务记录失败: {task_id}, 错误: {e}")
            return None
    
    async def update_task(self, task_record: TaskRecord) -> None:
        """更新任务记录"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            task_key = self.task_key_template.format(task_id=task_record.task_id)
            
            # 获取旧的任务记录（如果存在）
            old_value, _ = self.etcd_client.get(task_key)
            old_task_record = None
            if old_value:
                old_task_dict = json.loads(old_value.decode())
                old_task_record = self._dict_to_task_record(old_task_dict)
            
            # 转换为字典格式
            task_dict = self._task_record_to_dict(task_record)
            task_json = json.dumps(task_dict)
            
            # 创建租约
            lease = self.etcd_client.lease(self.default_ttl)
            
            # 使用事务进行原子操作
            def transaction():
                # 保存任务记录
                self.etcd_client.put(task_key, task_json, lease=lease)
                
                # 更新状态索引
                if old_task_record and old_task_record.status != task_record.status:
                    # 从旧状态中移除
                    old_status_key = self.task_status_key_template.format(
                        status=old_task_record.status.value
                    )
                    old_status_value, _ = self.etcd_client.get(old_status_key)
                    if old_status_value:
                        old_status_set = set(json.loads(old_status_value.decode()))
                        old_status_set.discard(task_record.task_id)
                        self.etcd_client.put(old_status_key, json.dumps(list(old_status_set)), lease=lease)
                
                # 添加到新状态
                status_key = self.task_status_key_template.format(status=task_record.status.value)
                status_value, _ = self.etcd_client.get(status_key)
                status_set = set(json.loads(status_value.decode())) if status_value else set()
                status_set.add(task_record.task_id)
                self.etcd_client.put(status_key, json.dumps(list(status_set)), lease=lease)
                
                # 更新服务器任务索引
                if task_record.assigned_server:
                    server_tasks_key = self.server_tasks_key_template.format(
                        server_id=task_record.assigned_server
                    )
                    server_value, _ = self.etcd_client.get(server_tasks_key)
                    server_set = set(json.loads(server_value.decode())) if server_value else set()
                    server_set.add(task_record.task_id)
                    self.etcd_client.put(server_tasks_key, json.dumps(list(server_set)), lease=lease)
                
                # 更新任务索引
                index_value, _ = self.etcd_client.get(self.task_index_key)
                index_dict = json.loads(index_value.decode()) if index_value else {}
                index_dict[task_record.task_id] = task_record.submit_time.timestamp()
                self.etcd_client.put(self.task_index_key, json.dumps(index_dict), lease=lease)
            
            # 执行事务
            transaction()
            
            self.logger.debug(f"已更新任务记录: {task_record.task_id}")
            
        except Exception as e:
            self.logger.error(f"更新任务记录失败: {task_record.task_id}, 错误: {e}")
            raise
    
    async def list_tasks(self, 
                        status: Optional[TaskStatus] = None,
                        server_id: Optional[str] = None,
                        limit: int = 100) -> List[TaskRecord]:
        """列出任务记录"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            task_ids = set()
            
            if status:
                # 从状态索引获取任务ID
                status_key = self.task_status_key_template.format(status=status.value)
                status_value, _ = self.etcd_client.get(status_key)
                if status_value:
                    status_set = set(json.loads(status_value.decode()))
                    task_ids.update(status_set)
            
            if server_id:
                # 从服务器索引获取任务ID
                server_tasks_key = self.server_tasks_key_template.format(server_id=server_id)
                server_value, _ = self.etcd_client.get(server_tasks_key)
                server_set = set()
                if server_value:
                    server_set = set(json.loads(server_value.decode()))
                
                if task_ids:
                    task_ids = task_ids.intersection(server_set)
                else:
                    task_ids = server_set
            
            if not task_ids and not status and not server_id:
                # 如果没有过滤条件，从任务索引获取最新的任务
                index_value, _ = self.etcd_client.get(self.task_index_key)
                if index_value:
                    index_dict = json.loads(index_value.decode())
                    # 按时间戳排序，获取最新的任务
                    sorted_tasks = sorted(
                        index_dict.items(), 
                        key=lambda x: x[1], 
                        reverse=True
                    )
                    task_ids = set(task_id for task_id, _ in sorted_tasks[:limit])
            
            # 限制结果数量
            task_ids = list(task_ids)[:limit]
            
            # 批量获取任务记录
            tasks = []
            for task_id in task_ids:
                task_record = await self.get_task(task_id)
                if task_record:
                    tasks.append(task_record)
            
            return tasks
            
        except Exception as e:
            self.logger.error(f"列出任务记录失败: {e}")
            return []
    
    async def delete_task(self, task_id: str) -> bool:
        """删除任务记录"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            # 获取任务记录以获取索引信息
            task_record = await self.get_task(task_id)
            if not task_record:
                return False
            
            # 使用事务进行原子删除
            def delete_transaction():
                # 删除任务记录
                task_key = self.task_key_template.format(task_id=task_id)
                self.etcd_client.delete(task_key)
                
                # 从状态索引中移除
                status_key = self.task_status_key_template.format(status=task_record.status.value)
                status_value, _ = self.etcd_client.get(status_key)
                if status_value:
                    status_set = set(json.loads(status_value.decode()))
                    status_set.discard(task_id)
                    if status_set:
                        self.etcd_client.put(status_key, json.dumps(list(status_set)))
                    else:
                        self.etcd_client.delete(status_key)
                
                # 从服务器索引中移除
                if task_record.assigned_server:
                    server_tasks_key = self.server_tasks_key_template.format(
                        server_id=task_record.assigned_server
                    )
                    server_value, _ = self.etcd_client.get(server_tasks_key)
                    if server_value:
                        server_set = set(json.loads(server_value.decode()))
                        server_set.discard(task_id)
                        if server_set:
                            self.etcd_client.put(server_tasks_key, json.dumps(list(server_set)))
                        else:
                            self.etcd_client.delete(server_tasks_key)
                
                # 从任务索引中移除
                index_value, _ = self.etcd_client.get(self.task_index_key)
                if index_value:
                    index_dict = json.loads(index_value.decode())
                    if task_id in index_dict:
                        del index_dict[task_id]
                        if index_dict:
                            self.etcd_client.put(self.task_index_key, json.dumps(index_dict))
                        else:
                            self.etcd_client.delete(self.task_index_key)
            
            # 执行删除事务
            delete_transaction()
            
            self.logger.info(f"已删除任务记录: {task_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"删除任务记录失败: {task_id}, 错误: {e}")
            return False
    
    async def get_task_count(self, status: Optional[TaskStatus] = None) -> int:
        """获取任务数量"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            if status:
                status_key = self.task_status_key_template.format(status=status.value)
                status_value, _ = self.etcd_client.get(status_key)
                if status_value:
                    status_set = set(json.loads(status_value.decode()))
                    return len(status_set)
                return 0
            else:
                index_value, _ = self.etcd_client.get(self.task_index_key)
                if index_value:
                    index_dict = json.loads(index_value.decode())
                    return len(index_dict)
                return 0
                
        except Exception as e:
            self.logger.error(f"获取任务数量失败: {e}")
            return 0
    
    async def get_server_task_count(self, server_id: str) -> int:
        """获取服务器任务数量"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            server_tasks_key = self.server_tasks_key_template.format(server_id=server_id)
            server_value, _ = self.etcd_client.get(server_tasks_key)
            if server_value:
                server_set = set(json.loads(server_value.decode()))
                return len(server_set)
            return 0
            
        except Exception as e:
            self.logger.error(f"获取服务器任务数量失败: {server_id}, 错误: {e}")
            return 0
    
    async def cleanup_expired_tasks(self, max_age_days: int = 30) -> int:
        """清理过期任务"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            cutoff_timestamp = (datetime.now(timezone.utc) - timedelta(days=max_age_days)).timestamp()
            
            # 获取任务索引
            index_value, _ = self.etcd_client.get(self.task_index_key)
            if not index_value:
                return 0
            
            index_dict = json.loads(index_value.decode())
            
            # 找出过期的任务
            expired_task_ids = [
                task_id for task_id, timestamp in index_dict.items()
                if timestamp < cutoff_timestamp
            ]
            
            # 批量删除过期任务
            deleted_count = 0
            for task_id in expired_task_ids:
                if await self.delete_task(task_id):
                    deleted_count += 1
            
            self.logger.info(f"已清理 {deleted_count} 个过期任务")
            return deleted_count
            
        except Exception as e:
            self.logger.error(f"清理过期任务失败: {e}")
            return 0
    
    async def acquire_lock(self, resource: str, timeout: int = 30) -> bool:
        """获取分布式锁"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            lock_key = self.lock_key_template.format(resource=resource)
            lock_value = f"{time.time()}_{id(self)}"
            
            # 创建租约
            lease = self.etcd_client.lease(timeout)
            
            # 尝试获取锁
            success, _ = self.etcd_client.transaction(
                compare=[
                    self.etcd_client.transactions.version(lock_key) == 0
                ],
                success=[
                    self.etcd_client.transactions.put(lock_key, lock_value, lease=lease)
                ],
                failure=[]
            )
            
            return success
            
        except Exception as e:
            self.logger.error(f"获取锁失败: {resource}, 错误: {e}")
            return False
    
    async def release_lock(self, resource: str) -> bool:
        """释放分布式锁"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            lock_key = self.lock_key_template.format(resource=resource)
            result = self.etcd_client.delete(lock_key)
            return result
            
        except Exception as e:
            self.logger.error(f"释放锁失败: {resource}, 错误: {e}")
            return False
    
    async def watch_tasks(self, 
                         status: Optional[TaskStatus] = None,
                         callback: Optional[callable] = None) -> None:
        """监听任务变化"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            if status:
                watch_key = self.task_status_key_template.format(status=status.value)
            else:
                watch_key = self.task_index_key
            
            # 创建监听器
            events_iterator, cancel = self.etcd_client.watch(watch_key)
            
            try:
                for event in events_iterator:
                    if callback:
                        await callback(event)
            finally:
                cancel()
                
        except Exception as e:
            self.logger.error(f"监听任务变化失败: {e}")
    
    async def get_statistics(self) -> Dict[str, Any]:
        """获取存储统计信息"""
        if not self.etcd_client:
            await self.connect()
        
        try:
            stats = {
                "total_tasks": await self.get_task_count(),
                "status_counts": {},
                "server_counts": {},
                "etcd_info": {}
            }
            
            # 获取各状态的任务数量
            for status in TaskStatus:
                count = await self.get_task_count(status)
                stats["status_counts"][status.value] = count
            
            # 获取etcd信息
            try:
                etcd_status = self.etcd_client.status()
                stats["etcd_info"] = {
                    "version": etcd_status.version,
                    "db_size": etcd_status.db_size,
                    "leader": etcd_status.leader,
                    "raft_index": etcd_status.raft_index,
                    "raft_term": etcd_status.raft_term
                }
            except Exception as e:
                self.logger.warning(f"获取etcd状态失败: {e}")
                stats["etcd_info"] = {"error": str(e)}
            
            return stats
            
        except Exception as e:
            self.logger.error(f"获取统计信息失败: {e}")
            return {}
    
    def _task_record_to_dict(self, task_record: TaskRecord) -> Dict[str, Any]:
        """将任务记录转换为字典"""
        return {
            "task_id": task_record.task_id,
            "client_ip": task_record.client_ip,
            "submit_time": task_record.submit_time.isoformat(),
            "assigned_server": task_record.assigned_server,
            "server_endpoint": task_record.server_endpoint,
            "status": task_record.status.value,
            "result_url": task_record.result_url,
            "processing_duration": task_record.processing_duration,
            "error_message": task_record.error_message,
            "correlation_id": task_record.correlation_id,
            "metadata": task_record.metadata or {}
        }
    
    def _dict_to_task_record(self, task_dict: Dict[str, Any]) -> TaskRecord:
        """将字典转换为任务记录"""
        return TaskRecord(
            task_id=task_dict["task_id"],
            client_ip=task_dict["client_ip"],
            submit_time=datetime.fromisoformat(task_dict["submit_time"]),
            assigned_server=task_dict.get("assigned_server"),
            server_endpoint=task_dict.get("server_endpoint"),
            status=TaskStatus(task_dict["status"]),
            result_url=task_dict.get("result_url"),
            processing_duration=task_dict.get("processing_duration"),
            error_message=task_dict.get("error_message"),
            correlation_id=task_dict.get("correlation_id"),
            metadata=task_dict.get("metadata", {})
        )
